import os
import pickle
from typing import TYPE_CHECKING
from MelodieFuncFlow import *

if TYPE_CHECKING:
    from .indexing import DocInfo

_data_folder = ".py-miner-web"


def set_data_folder(folder: str):
    global _data_folder
    _data_folder = folder
    if not os.path.exists(_data_folder):
        os.makedirs(_data_folder)


def data_storage_path(*paths):
    path = os.path.join(_data_folder, *paths)
    if not os.path.exists(path):
        os.makedirs(path)
    return path


def file_indexed(md5: str):
    path = os.path.join(doc_storage_path(), md5 + ".pkl")
    return os.path.exists(path)


doc_storage_path = lambda: data_storage_path("parsed_docs")
config_storage_path = lambda: data_storage_path("config")


def dump_document(doc: "DocInfo"):
    path = os.path.join(doc_storage_path(), doc.hash + ".pkl")
    with open(path, "wb") as f:
        pickle.dump(doc, f)


def iter_documents_info() -> MelodieGenerator["DocInfo"]:
    def _():
        for item in os.listdir(doc_storage_path()):
            with open(os.path.join(doc_storage_path(), item), "rb") as f:
                yield pickle.load(f)

    return MelodieGenerator(_())


def iter_documents_in_folder(folder_to_index: str) -> MelodieGenerator[str]:
    def _():
        for root, dirs, files in os.walk(folder_to_index):
            for file in files:
                yield os.path.abspath(os.path.join(root, file))

    return MelodieGenerator(_())
